import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
from tqdm import tqdm
import torch
import torch.nn as nn
import numpy as np
import math
import pandas as pd

from sklearn.model_selection import train_test_split
from sklearn.preprocessing import MinMaxScaler

from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error, r2_score

import xgboost as xgb


def round_and_none_neg(n):
    if n < 0:
        return 0
    else:
        return round(n)


## 载入模型d
model = xgb.Booster()
model.load_model("model.bin")

predict_data = pd.read_csv("predict_data.csv").drop(columns="Unnamed: 0")
features = predict_data.iloc[:, :]
dtest = xgb.DMatrix(features)

predicts = model.predict(dtest)

weibo_predict_data = pd.read_csv(
    "./data/weibo_predict_data.txt",
    sep="\t",
    header=None,
    names=["uid", "mid", "time", "content"],
)
file_out = []

for index, row in tqdm(weibo_predict_data.iterrows(), total=len(weibo_predict_data)):
    # file_out.append("{0}\t{1}\t{2},{3},{4}\n".format(row["uid"],row["mid"],predicts[index][0], predicts[index][1], predicts[index][2]))

    file_out.append(
        "{0}\t{1}\t{2},{3},{4}\n".format(
            row["uid"],
            row["mid"],
            max(round_and_none_neg(predicts[index][0]) - 1, 0),
            max(round_and_none_neg(predicts[index][1]) - 1, 0),
            max(round_and_none_neg(predicts[index][2]) - 1, 0),
        )
    )

with open("answer_xgboost.txt", "w", encoding="utf-8") as f:
    f.writelines(file_out)
